#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
@FileName: predict.py
@Author: Hugo Wang
@Date: 2025-06-06 16:40
@Project: company_talent_loss
"""
import pandas as pd
import numpy as np
import os
import datetime
from utils.log import Logger
from sklearn.preprocessing import StandardScaler
from sklearn.preprocessing import LabelEncoder
import joblib
from sklearn.metrics import accuracy_score, roc_auc_score
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt
plt.rcParams['font.family'] = 'SimHei'
plt.rcParams['font.size'] = 15

class PredictResult:
    def __init__(self,file_path):
        # 1、配置日志
        logfileName = 'predict_' + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
        self.logger = Logger("../", logfileName).get_logger()
        # 2、获取数据源
        self.data_source = pd.read_csv(file_path)

    def pred_feature_extract(self, test_data, logger):
        #数据基本处理
        df = test_data[
            ['Age', 'Department', 'DistanceFromHome', 'Education', 'EnvironmentSatisfaction', 'Gender', 'JobLevel',
             'JobSatisfaction', 'MaritalStatus', 'MonthlyIncome', 'NumCompaniesWorked', 'OverTime', 'PercentSalaryHike',
             'PerformanceRating', 'RelationshipSatisfaction', 'StockOptionLevel', 'TotalWorkingYears',
             'WorkLifeBalance', 'YearsAtCompany', 'YearsInCurrentRole', 'YearsSinceLastPromotion',
             'YearsWithCurrManager', 'Attrition']]
        #对特征列进行值的替换处理
        df['MaritalStatus'] = df['MaritalStatus'].apply(lambda x: 0 if x in ['Divorced', 'Single'] else 1)
        df['Department'] = df['Department'].apply(lambda x: 0 if x in ['Sales', 'Human Resources'] else 1)
        #对特征中类型为Object的特征进行热编码处理
        df2 = pd.get_dummies(df)
        #相似特征列的删除
        df3 = df2.drop(['Gender_Male', 'OverTime_No'], axis=1)
        # print(df3.info())
        # 从预处理后的数据列中找出特征列和标签
        x = df3[['Age', 'EnvironmentSatisfaction', 'JobSatisfaction', 'MaritalStatus', 'MonthlyIncome',
                 'NumCompaniesWorked', 'PercentSalaryHike', 'PerformanceRating', 'RelationshipSatisfaction',
                 'StockOptionLevel', 'TotalWorkingYears', 'WorkLifeBalance', 'YearsAtCompany', 'YearsInCurrentRole',
                 'YearsSinceLastPromotion', 'Gender_Female', 'OverTime_Yes']]
        y = df3['Attrition']
        return x,y

if __name__ == '__main__':
    pr = PredictResult('../data/test2.csv')
    #获取测试数据
    test_data = pr.data_source
    #获取特征和标签
    x,y = pr.pred_feature_extract(test_data,pr.logger)

    # 对标签进行编码处理
    le = LabelEncoder()
    y = le.fit_transform(y)
    # 对特征进行标准化处理
    transfer = StandardScaler()
    x = transfer.fit_transform(x)

    # 加载本地的模型
    # dtcModel = joblib.load('../model/2025060612_dtcModel.pkl')
    dtcModel = joblib.load('../model/2025060612_xgbModel.pkl')
    y_pred = dtcModel.predict(x)
    # print(y_pred.to_string())
    print(f"测试集的准确率：{accuracy_score(y, y_pred)}")
    print(f"测试集的AUC值{roc_auc_score(y, y_pred)}")



